package rabbit.flt.plugins.reactor.plugin;

import rabbit.flt.common.context.TraceContext;
import rabbit.flt.common.context.TraceMetaData;
import rabbit.flt.common.utils.ReflectUtils;
import rabbit.flt.plugins.common.plugin.SupportPlugin;

import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.*;

public class ReactorExchangeSupportPlugin extends SupportPlugin {

    private Map<Predicate, Function<Object, Object>> handleMap = new ConcurrentHashMap<>();

    public ReactorExchangeSupportPlugin() {
        handleMap.put(obj -> obj instanceof Function, input -> createExchangeFunction((Function) input));
        handleMap.put(obj -> obj instanceof BiConsumer, input -> createExchangeBiConsumer((BiConsumer) input));
        handleMap.put(obj -> obj instanceof Runnable, input -> getRunnable((Runnable) input));
        handleMap.put(obj -> obj instanceof Callable, input -> getCallable((Callable) input));
        handleMap.put(obj -> obj instanceof Supplier, input -> getSupplier((Supplier) input));
        handleMap.put(obj -> obj instanceof Consumer, input -> createExchangeConsumer((Consumer) input));
    }

    /**
     * 增强map 和 flatmap
     *
     * @param target
     * @param method
     * @param args
     * @return
     */
    @Override
    public Object[] before(Object target, Method method, Object[] args) {
        if (isTraceOpened()) {
            for (int i = 0; i < args.length; i++) {
                for (Map.Entry<Predicate, Function<Object, Object>> entry : handleMap.entrySet()) {
                    if (entry.getKey().test(args[i])) {
                        args[i] = entry.getValue().apply(args[i]);
                    }
                }
            }
        }
        return super.before(target, method, args);
    }

    protected Function createExchangeFunction(Function mapFunc) {
        TraceMetaData metaData = TraceContext.getTraceMetaData();
        String rootSpanId = TraceContext.getRootSpanId();
        return t -> {
            boolean opened = tryOpenTrace(metaData, rootSpanId);
            try {
                return mapFunc.apply(t);
            } finally {
                if (opened) {
                    TraceContext.clearContext();
                }
            }
        };
    }

    protected BiConsumer createExchangeBiConsumer(BiConsumer biConsumer) {
        TraceMetaData metaData = TraceContext.getTraceMetaData();
        String rootSpanId = TraceContext.getRootSpanId();
        return (t, u) -> {
            boolean opened = tryOpenTrace(metaData, rootSpanId);
            try {
                biConsumer.accept(t, u);
            } finally {
                if (opened) {
                    TraceContext.clearContext();
                }
            }
        };
    }

    protected Consumer createExchangeConsumer(Consumer consumer) {
        TraceMetaData metaData = TraceContext.getTraceMetaData();
        String rootSpanId = TraceContext.getRootSpanId();
        return t -> {
            boolean opened = tryOpenTrace(metaData, rootSpanId);
            try {
                consumer.accept(t);
            } finally {
                if (opened) {
                    TraceContext.clearContext();
                }
            }
        };
    }

    protected Supplier getSupplier(Supplier supplier) {
        TraceMetaData metaData = TraceContext.getTraceMetaData();
        String rootSpanId = TraceContext.getRootSpanId();
        return () -> {
            boolean opened = tryOpenTrace(metaData, rootSpanId);
            try {
                return supplier.get();
            } finally {
                if (opened) {
                    TraceContext.clearContext();
                }
            }
        };
    }

    protected Callable getCallable(Callable callable) {
        TraceMetaData metaData = TraceContext.getTraceMetaData();
        String rootSpanId = TraceContext.getRootSpanId();
        return () -> {
            boolean opened = tryOpenTrace(metaData, rootSpanId);
            try {
                return ReflectUtils.call((Callable) callable::call);
            } finally {
                if (opened) {
                    TraceContext.clearContext();
                }
            }
        };
    }

    protected Runnable getRunnable(Runnable runnable) {
        TraceMetaData metaData = TraceContext.getTraceMetaData();
        String rootSpanId = TraceContext.getRootSpanId();
        return () -> {
            boolean opened = tryOpenTrace(metaData, rootSpanId);
            try {
                runnable.run();
            } finally {
                if (opened) {
                    TraceContext.clearContext();
                }
            }
        };
    }

    /**
     * 尝试开启trace
     *
     * @param metaData
     * @param rootSpanId
     * @return
     */
    private final boolean tryOpenTrace(TraceMetaData metaData, String rootSpanId) {
        if (!isTraceOpened()) {
            // 传递值
            TraceContext.openTraceByMeta(metaData);
            TraceContext.initRootSpanId(rootSpanId);
            return true;
        }
        return false;
    }

}
